package com.litongjava.ai.server.padddle.ocr.service.impl;

import com.alibaba.fastjson.JSONObject;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.litongjava.ai.server.padddle.ocr.common.UrlConstants;
import com.litongjava.ai.server.padddle.ocr.model.ChatMessage;
import com.litongjava.ai.server.padddle.ocr.model.ChatRequest;
import com.litongjava.ai.server.padddle.ocr.model.Options;
import com.litongjava.ai.server.padddle.ocr.service.ChatService;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

import static com.litongjava.ai.server.padddle.ocr.common.UrlConstants.MODEL;

@Slf4j
@Service
public class ChatServiceImpl implements ChatService {

    private final OkHttpClient client;
    private final ObjectMapper objectMapper;

    public ChatServiceImpl() {
        this.client = new OkHttpClient.Builder()
                .connectTimeout(30, TimeUnit.SECONDS)
                .readTimeout(0, TimeUnit.SECONDS)  // 无限读取超时
                .writeTimeout(30, TimeUnit.SECONDS)
                .build();
        this.objectMapper = new ObjectMapper();
    }

    @Override
    public StreamingResponseBody streamChat(ChatRequest request) {
        return outputStream -> {
            try {
                // 构建请求体
                Map<String, Object> requestBody = new HashMap<>();
                requestBody.put("model", MODEL);
                
                if (request.isUseContext() && request.getMessages() != null && !request.getMessages().isEmpty()) {
                    // 使用上下文模式
                    requestBody.put("messages", request.getMessages());
                    // 添加当前用户的新消息
                    if (request.getPrompt() != null && !request.getPrompt().isEmpty()) {
                        request.getMessages().add(new ChatMessage("user", request.getPrompt()));
                    }
                } else {
                    // 不使用上下文，仅发送当前prompt
                    requestBody.put("prompt", request.getPrompt());
                }
                
                String jsonBody = objectMapper.writeValueAsString(requestBody);
                
                MediaType JSON = MediaType.parse("application/json; charset=utf-8");
                RequestBody body = RequestBody.create(jsonBody, JSON);
                
                Request httpRequest = new Request.Builder()
                        .url(UrlConstants.DEEP_URL)
                        .post(body)
                        .build();
                
                log.info("发送请求到Ollama: {}", jsonBody);
                
                try (Response response = client.newCall(httpRequest).execute()) {
                    if (!response.isSuccessful()) {
                        throw new IOException("请求失败: " + response);
                    }
                    
                    ResponseBody responseBody = response.body();
                    if (responseBody == null) {
                        throw new IOException("响应体为空");
                    }
                    
                    try (BufferedReader reader = new BufferedReader(
                            new InputStreamReader(responseBody.byteStream(), StandardCharsets.UTF_8))) {
                        
                        StringBuilder responseContent = new StringBuilder();
                        String line;
                        while ((line = reader.readLine()) != null) {
                            try {
                                // 解析JSON响应
                                JSONObject dataJson = JSONObject.parseObject(line);
                                if (dataJson.containsKey("message")) {
                                    JSONObject textJson = dataJson.getJSONObject("message");
                                    if(textJson.containsKey("content")) {
                                        String text = textJson.getString("content");
                                        // 过滤掉特殊标记和空行
                                        if (text != null && !text.equals("<think>")
                                                && !text.equals("</think>")
                                                && !text.startsWith("\n")) {
                                            outputStream.write(text.getBytes(StandardCharsets.UTF_8));
                                            outputStream.flush();
                                            responseContent.append(text);
                                        }
                                    }
                                }
                                
                                // 检查是否完成
                                if (dataJson.containsKey("done") && dataJson.getBoolean("done")) {
                                    // 如果使用上下文，将AI的回复添加到消息历史
                                    if (request.isUseContext() && request.getMessages() != null) {
                                        request.getMessages().add(new ChatMessage("assistant", responseContent.toString()));
                                    }
                                    break;
                                }
                            } catch (Exception e) {
                                log.error("解析JSON发生错误: {}", line, e);
                            }
                        }
                    }
                }
            } catch (Exception e) {
                log.error("处理流式响应时发生错误", e);
                try {
                    String errorMsg = "发生错误: " + e.getMessage();
                    outputStream.write(errorMsg.getBytes(StandardCharsets.UTF_8));
                } catch (IOException ioException) {
                    log.error("写入错误消息时发生异常", ioException);
                }
            }
        };
    }

    @Override
    public StreamingResponseBody streamChat(List<ChatMessage> request) {
        return outputStream -> {
            try {
                Options options = new Options();
                // 构建请求体
                Map<String, Object> requestBody = new HashMap<>();
                requestBody.put("model", MODEL);
                requestBody.put("messages", request);
                requestBody.put("options",options);
                String jsonBody = objectMapper.writeValueAsString(requestBody);

                MediaType JSON = MediaType.parse("application/json; charset=utf-8");
                RequestBody body = RequestBody.create(jsonBody, JSON);

                Request httpRequest = new Request.Builder()
                        .url(UrlConstants.DEEP_URL)
                        .post(body)
                        .build();


                try (Response response = client.newCall(httpRequest).execute()) {
                    if (!response.isSuccessful()) {
                        throw new IOException("请求失败: " + response);
                    }

                    ResponseBody responseBody = response.body();
                    if (responseBody == null) {
                        throw new IOException("响应体为空");
                    }

                    try (BufferedReader reader = new BufferedReader(
                            new InputStreamReader(responseBody.byteStream(), StandardCharsets.UTF_8))) {

                        String line;
                        while ((line = reader.readLine()) != null) {
                            try {
                                // 解析JSON响应
                                JSONObject dataJson = JSONObject.parseObject(line);
                                if (dataJson.containsKey("message")) {
                                    JSONObject textJson = dataJson.getJSONObject("message");
                                    if(textJson.containsKey("content")) {
                                        String text = textJson.getString("content");
                                        // 过滤掉特殊标记和空行
                                        if (text != null && !text.equals("<think>")
                                                && !text.equals("</think>")
                                                && !text.startsWith("\n")) {
                                            outputStream.write(text.getBytes(StandardCharsets.UTF_8));
                                            outputStream.flush();
                                        }
                                    }
                                }

                                // 检查是否完成
                                if (dataJson.containsKey("done") && dataJson.getBoolean("done")) {
                                    break;
                                }
                            } catch (Exception e) {
                                log.error("解析JSON发生错误: {}", line, e);
                            }
                        }
                    }
                }
            } catch (Exception e) {
                log.error("处理流式响应时发生错误", e);
                try {
                    String errorMsg = "发生错误: " + e.getMessage();
                    outputStream.write(errorMsg.getBytes(StandardCharsets.UTF_8));
                } catch (IOException ioException) {
                    log.error("写入错误消息时发生异常", ioException);
                }
            }
        };
    }
} 